import pickle
import json
import ast

def load_topic_reports(pkl_path, max_pro=5, max_con=5):
    with open(pkl_path, "rb") as f:
        data = pickle.load(f)

    topic_reports = {}

    for topic, views in data.items():
        pro_parts = []
        con_parts = []

        for view in views:
            if not isinstance(view, list) or len(view) < 3:
                continue

            stance, subtitle, content = view
            tag = None
            if "pro" in stance.lower() and len(pro_parts) < max_pro:
                tag = "【支持观点】"
                pro_parts.append(f"{tag}（{subtitle}）\n{content}")
            elif "con" in stance.lower() and len(con_parts) < max_con:
                tag = "【反对观点】"
                con_parts.append(f"{tag}（{subtitle}）\n{content}")

            # 如果两个都够了就提前退出
            if len(pro_parts) >= max_pro and len(con_parts) >= max_con:
                break

        combined = pro_parts + con_parts
        if combined:  # 至少有一边不为空才加入
            topic_reports[topic] = "\n\n".join(combined)

    return topic_reports

def load_jsonl(file_path):
    profiles = []
    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    profiles.append(ast.literal_eval(line))  # 处理字符串
                except Exception as e:
                    print(f"[!] Failed parsing line: {line}\nError: {e}")
    return profiles

def load_full_profiles(personality_path, profile_path):
    personalities = load_jsonl(personality_path)
    backgrounds = load_jsonl(profile_path)
    full_profiles = []

    for p in personalities:
        for b in backgrounds:
            full_profiles.append(p + " " + b)  # 合并为完整 profile 描述

    return full_profiles

if __name__ == "__main__":
    # 示例用法
    topic_reports = load_topic_reports("/data01/public/chenmiaokun/everything/cognitive_manipulation/datasets/opinion_pro/procons.pkl")
    print(f"加载了 {len(topic_reports)} 个话题报告")

    # 加载个性和背景数据
    personalities = load_jsonl("/data01/public/chenmiaokun/everything/cognitive_manipulation/datasets/profile/personality.jsonl")
    backgrounds = load_jsonl("/data01/public/chenmiaokun/everything/cognitive_manipulation/datasets/profile/profile.jsonl")

    # 生成完整的 profile 列表
    full_profiles = load_full_profiles("/data01/public/chenmiaokun/everything/cognitive_manipulation/datasets/profile/personality.jsonl", "/data01/public/chenmiaokun/everything/cognitive_manipulation/datasets/profile/profile.jsonl")
    print(f"生成了 {len(full_profiles)} 个完整的 profile 描述")


